import datetime
import os
import time

import torch
import torch.utils.data
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor

from data.apple_dataset import AppleDataset
from utility.engine import train_one_epoch, evaluate

import utility.utils as utils
import utility.transforms as T

######################################################
# Train either a Faster-RCNN or Mask-RCNN predictor
# using the MinneApple dataset
######################################################


# BEGIN
def get_transform(train):
    """
    生成用于数据转换的函数

    参数:
    train (bool): 指示是否为训练模式的布尔值

    返回:
    torchvision.transforms.Compose: 一个包含所有数据转换操作的组合对象
    """
    
    transforms = []
    transforms.append(T.ToTensor())  # 将图像转换为张量
    if train:
        transforms.append(T.RandomHorizontalFlip(0.5))  # 以0.5的概率水平翻转图像
    return T.Compose(transforms)  # 返回所有数据转换操作的组合对象
# END


# BEGIN
def get_maskrcnn_model_instance(num_classes):
    # 加载在COCO数据集上预训练的实例分割模型
    model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)

    # 获取分类器的输入特征数
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    # 用新的分类器替换预训练的分类器
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    # 获取掩膜分类器的输入特征数
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer = 256
    # 用新的掩膜分类器替换预训练的分类器
    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, hidden_layer, num_classes)
    return model
# END


# 定义一个函数，用于获取一个 Faster R-CNN 模型实例，参数为 num_classes 表示类别数量
def get_frcnn_model_instance(num_classes):
    # 载入一个在 COCO 数据集上预训练的实例分割模型
    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)

    # 获取分类器的输入特征数量
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    # 用一个新的分类器替换预训练的分类器
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    return model


# BEGIN
# 定义主函数，接收参数args
def main(args):
    # 打印参数args
    print(args)
    # 获取设备信息
    device = args.device

    # 数据加载代码
    print("Loading data")
    num_classes = 2
    # 加载苹果数据集，传入训练路径和变换函数
    dataset = AppleDataset(os.path.join(args.data_path, 'train'), get_transform(train=True))
    # 加载测试数据集，传入测试路径和变换函数
    dataset_test = AppleDataset(os.path.join(args.data_path, 'test'), get_transform(train=False))

    print("Creating data loaders")
    # 创建训练数据加载器，设置批量大小、是否打乱数据、工作进程数和数据整理函数
    data_loader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=True,
                                              num_workers=args.workers, collate_fn=utils.collate_fn)

    # 创建测试数据加载器，设置批量大小为1、不打乱数据、工作进程数和数据整理函数
    data_loader_test = torch.utils.data.DataLoader(dataset_test, batch_size=1,
                                                   shuffle=False, num_workers=args.workers,
                                                   collate_fn=utils.collate_fn)

    print("Creating model")
    # 创建模型，根据参数args中的模型类型选择不同的模型
    if args.model == 'maskrcnn':
        model = get_maskrcnn_model_instance(num_classes)
    else:
 

        # 获取 Faster R-CNN 模型实例
        model = get_frcnn_model_instance(num_classes)

    # 将模型移动到正确的设备上
    model.to(device)

    # 获取所有需要梯度更新的参数
    params = [p for p in model.parameters() if p.requires_grad]
    
    # 使用 SGD 优化器进行模型参数更新
    optimizer = torch.optim.SGD(params, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)

    # 使用 MultiStepLR 学习率调度器
    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.lr_steps, gamma=args.lr_gamma)

    # 如果需要从之前的训练中恢复，则加载模型、优化器和学习率调度器状态
    if args.resume:
        checkpoint = torch.load(args.resume, map_location='cpu')
        model.load_state_dict(checkpoint['model'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])

    # 开始训练过程
    print("Start training")
    start_time = time.time()
    
    # 遍历每个 epoch 进行训练
    for epoch in range(args.epochs):
        train_one_epoch(model, optimizer, data_loader, device, epoch, args.print_freq)
        
        # 更新学习率
        lr_scheduler.step()

    if args.output_dir:  # 检查是否设置了输出目录
        torch.save({  # 使用torch.save保存模型和优化器状态
        'epoch': epoch,  # 保存当前训练的轮数
        'model': model.state_dict(),  # 保存模型的状态字典
        'optimizer': optimizer.state_dict(),  # 保存优化器的状态字典
        'lr_scheduler': lr_scheduler.state_dict(),  # 保存学习率调度器的状态字典
        },  os.path.join(args.output_dir, 'model_{}.pth'.format(epoch)))  # 将训练好的模型保存到指定的输出目录中

    # evaluate after every epoch
    evaluate(model, data_loader_test, device=device)  # 每个epoch结束后对模型进行评估

total_time = time.time() - start_time  # 计算训练总时间
total_time_str = str(datetime.timedelta(seconds=int(total_time)))  # 将总时间转换为可读的时间格式
print('Training time {}'.format(total_time_str))  # 打印训练时间信息


# BEGIN

# 如果作为独立脚本运行
if __name__ == "__main__":
    import argparse
    # 创建参数解析器
    parser = argparse.ArgumentParser(description='PyTorch Detection Training')
    
    # 添加参数：数据集路径
    parser.add_argument('--data_path', default='/home/nicolai/phd/data/cvppp_2019/apple_dataset', help='dataset')
    
    # 添加参数：数据集名称
    parser.add_argument('--dataset', default='AppleDataset', help='dataset')
    
    # 添加参数：模型名称
    parser.add_argument('--model', default='maskrcnn', help='model')
    
    # 添加参数：设备类型
    parser.add_argument('--device', default='cuda', help='device')
    
    # 添加参数：批量大小
    parser.add_argument('-b', '--batch-size', default=2, type=int)
    
    # 添加参数：训练轮数
    parser.add_argument('--epochs', default=13, type=int, metavar='N', help='number of total epochs to run')
    
    # 添加参数：数据加载工作线程数
    parser.add_argument('-j', '--workers', default=4, type=int, metavar='N', help='number of data loading workers (default: 16)')
    
    # 添加参数：学习率
    parser.add_argument('--lr', default=0.02, type=float, help='initial learning rate')
    
    # 添加参数：动量
    parser.add_argument('--momentum', default=0.9, type=float, metavar='M', help='momentum')

    # 添加参数'--wd'或'--weight-decay'，默认值为1e-4，类型为浮点数，帮助信息为'weight decay (default: 1e-4)'，目标为'weight_decay'
    parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float, metavar='W', help='weight decay (default: 1e-4)', dest='weight_decay')

    # 添加参数'--lr-step-size'，默认值为8，类型为整数，帮助信息为'decrease lr every step-size epochs'
    parser.add_argument('--lr-step-size', default=8, type=int, help='decrease lr every step-size epochs')

    # 添加参数'--lr-steps'，默认值为[8, 11]，类型为整数列表，帮助信息为'decrease lr every step-size epochs'
    parser.add_argument('--lr-steps', default=[8, 11], nargs='+', type=int, help='decrease lr every step-size epochs')

    # 添加参数'--lr-gamma'，默认值为0.1，类型为浮点数，帮助信息为'decrease lr by a factor of lr-gamma'
    parser.add_argument('--lr-gamma', default=0.1, type=float, help='decrease lr by a factor of lr-gamma')

    # 添加参数'--print-freq'，默认值为20，类型为整数，帮助信息为'print frequency'
    parser.add_argument('--print-freq', default=20, type=int, help='print frequency')

    # 添加参数'--output-dir'，默认值为'.'，帮助信息为'path where to save'
    parser.add_argument('--output-dir', default='.', help='path where to save')

    # 添加参数'--resume'，默认值为空字符串，帮助信息为'resume from checkpoint'
    parser.add_argument('--resume', default='', help='resume from checkpoint')

    # 解析命令行参数
    args = parser.parse_args()

    # 打印模型参数
    print(args.model)

    # 断言模型参数为'mrcnn'或'frcnn'
    assert(args.model in ['mrcnn', 'frcnn'])

    # 如果有输出目录参数，则创建目录
    if args.output_dir:
        utils.mkdir(args.output_dir)

    # 调用主函数并传入参数
    main(args)
